from fastapi import APIRouter, Query
import baostock as bs
import pandas as pd

router = APIRouter()


@router.get("/data/SaveData")
async def create_user():
    #### 登陆系统 ####
    lg = bs.login()
    print('login respond error_code:' + lg.error_code)
    print('login respond  error_msg:' + lg.error_msg)

    #### 获取沪深A股历史K线数据 ####
    rs = bs.query_history_k_data_plus("sh.600000",
                                      "date,code,open,high,low,close,preclose,volume,amount,adjustflag,turn,tradestatus,pctChg,isST",
                                      start_date='2022-07-01', end_date='2022-12-31',
                                      frequency="d", adjustflag="3")
    print('query_history_k_data_plus respond error_code:' + rs.error_code)
    print('query_history_k_data_plus respond  error_msg:' + rs.error_msg)

    #### 打印结果集 ####
    data_list = []
    while (rs.error_code == '0') & rs.next():
        data_list.append(rs.get_row_data())
    result = pd.DataFrame(data_list, columns=rs.fields)

    #### 登出系统 ####
    bs.logout()

    # 将结果保存为 CSV 文件
    result.to_csv('stock_data.csv', index=False)

    print("数据已保存到 stock_data.csv 文件中")
    # 这里可以编写创建用户的逻辑
    return {"message": "数据已保存到 stock_data.csv 文件中"}


@router.get("/data/GetData")
async def get_stock_data(
        start_date: str = Query(None),
        end_date: str = Query(None)
):
    lg = bs.login()
    rs = bs.query_history_k_data_plus("sh.600000",
                                      "date,code,open,high,low,close,preclose,volume,amount,adjustflag,turn,"
                                      "tradestatus, "
                                      "pctChg,isST",
                                      start_date=start_date, end_date=end_date,
                                      frequency="d", adjustflag="3")
    data_list = []
    while (rs.error_code == '0') & rs.next():
        data_list.append(rs.get_row_data())
    result = pd.DataFrame(data_list, columns=rs.fields)
    bs.logout()

    if start_date and end_date:
        result = result[(result['date'] >= start_date) & (result['date'] <= end_date)]

    return result.to_dict(orient='records')

